# prediction_gui.py
import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QComboBox, QPushButton, QLabel,
                           QFileDialog, QMessageBox, QDialog, QTableWidget, QTableWidgetItem, QGridLayout)
from PyQt5.QtCore import Qt
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import pandas as pd
from prediction import Prediction, StockDataset
from prediction import CNNModel  # 确保导入 CNNModel 类
import torch
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import matplotlib.font_manager as fm
import os
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import pandas as pd


# 在 predict 方法中，绘制图表之前添加以下代码
font_path = r"E:\simfang.ttf"  # 字体路径
prop = fm.FontProperties(fname=font_path)
# 设置字体
font = FontProperties(fname=r"E:\simfang.ttf", size=14)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

class PredictionGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('股票预测系统')
        self.setGeometry(100, 100, 1200, 800)

        # 初始化预测类
        self.prediction = Prediction(data_days=10, batch_size=200)
        self.dataset = StockDataset(data_days=10, remake_data=False)
        
        # 读取股票代码和名字对应表
        stock_names_df = pd.read_csv('data/stock_names.csv', header=None, names=['code', 'name'])

        # 确保股票代码是纯字符串，不带前缀和空格
        stock_names_df['code'] = stock_names_df['code'].astype(str).str.strip().str[-6:]

        # 转换为字典
        self.stock_names = stock_names_df.set_index('code')['name'].to_dict()

        # 创建主窗口部件和布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)

        # 创建控制面板
        control_panel = QWidget()
        control_layout = QHBoxLayout(control_panel)
        
        # 添加模型选择下拉框
        self.model_combo = QComboBox()
        self.model_combo.addItems([ 'CNN', 'LSTM', 'RNN_RELU', 'ResNet18', 'ResNet34', 'ResNet50', 'ResNet101', 'ResNet152', 'DenseNet' ])
        control_layout.addWidget(QLabel('选择模型:'))
        control_layout.addWidget(self.model_combo)

        # 添加股票代码选择下拉框
        self.stock_combo = QComboBox()
        for code in self.dataset.stocks_codes[:20]:
            formatted_code = str(code).strip()[-6:]  # 取股票代码最后6位
            name = self.stock_names.get(formatted_code, '未知股票')
            self.stock_combo.addItem(f'{formatted_code} - {name}', formatted_code)

        control_layout.addWidget(QLabel('选择股票:'))
        control_layout.addWidget(self.stock_combo)

        # 添加预测按钮
        predict_btn = QPushButton('开始预测')
        predict_btn.clicked.connect(self.predict)
        control_layout.addWidget(predict_btn)
        # 在控制面板中添加模型比较按钮
        compare_btn = QPushButton('比较模型')
        compare_btn.clicked.connect(self.open_model_comparison)
        control_layout.addWidget(compare_btn)
        # 添加控制面板到主布局
        layout.addWidget(control_panel)

        # 创建图表
        self.figure = plt.figure(figsize=(12, 6))
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)
    def open_model_comparison(self):
        dialog = ModelComparisonDialog(self)
        dialog.exec_()
    def predict(self):
        try:
            # 清除之前的图表
            self.figure.clear()

            # 获取选择的模型和股票代码
            model_name = self.model_combo.currentText()
            stock_code = self.stock_combo.currentData()  # 获取股票代码

            # 股票代码格式处理
            pure_code = str(stock_code).strip()[-6:]  # 提取6位数字代码
            
            # 尝试多种可能的文件命名格式
            possible_paths = [
                f'./data/stocks/{pure_code}.csv',
                f'./data/stocks/sh.{pure_code}.csv',
                f'./data/stocks/sz.{pure_code}.csv'
            ]
            
            df = None
            actual_path = None
            
            for path in possible_paths:
                if os.path.exists(path):
                    try:
                        df = pd.read_csv(path)
                        actual_path = path
                        print(f"成功读取文件: {path}")
                        break
                    except Exception as e:
                        print(f"尝试读取 {path} 失败: {str(e)}")
            
            if df is None:
                raise FileNotFoundError(f"找不到股票代码为 {pure_code} 的数据文件，已尝试路径: {possible_paths}")

            # 确定用于预测的股票代码格式
            prediction_stock_code = pure_code
            if 'sh.' in actual_path:
                prediction_stock_code = f'sh.{pure_code}'
            elif 'sz.' in actual_path:
                prediction_stock_code = f'sz.{pure_code}'
                
            # 获取股票名字
            stock_name = self.stock_names.get(pure_code, '未知股票')
            
            # 清理列名
            df.columns = df.columns.str.strip().str.lower()
            print("列名检查:", df.columns)

            # 确保必要的列存在
            if 'date' not in df.columns:
                raise ValueError("数据文件中缺少 'date' 列，请检查文件格式。")
            if 'close' not in df.columns:
                raise ValueError("数据文件中缺少 'close' 列，请检查文件格式。")

            # 提取交易日期和收盘价
            trading_dates = df['date'].values
            true_close = df['close'].values
            x_r = range(0, len(trading_dates))

            # 定义收盘价预测函数
            def close_p(x):
                if type(x) == int:
                    return x
                x = x[0, 1].item()
                return x if 0.2 > x > -0.2 else 0.0

            # 根据选择的模型进行预测
            predict_func = getattr(self.prediction, f'predict_{model_name.lower()}')
            predict_close = [true_close[j] * (1 + close_p(predict_func(prediction_stock_code, (0, j))))
                            for j in range(len(trading_dates))]

            # 绘制图表
            ax = self.figure.add_subplot(111)
            ax.plot(x_r, true_close, label='真实值')
            ax.plot(x_r, predict_close, label=f'{model_name}模型预测值')
            
            # 设置x轴标签
            x_ticks = list(x_r[::100])
            if x_r[-1] not in x_ticks:
                x_ticks.append(x_r[-1])
            x_labels = [trading_dates[i] for i in x_ticks]
            ax.set_xticks(x_ticks)
            ax.set_xticklabels(x_labels, rotation=45, fontproperties=prop)

            ax.set_ylabel('开盘价', fontproperties=prop)
            ax.legend(prop=prop)
            
            # 添加股票名字标签
            ax.set_title(f'{stock_name} ({pure_code}) - {model_name}模型预测', fontproperties=prop)
            
            # 更新画布
            self.canvas.draw()

        except Exception as e:
            import traceback
            traceback.print_exc()  # 打印完整错误堆栈
            QMessageBox.critical(self, '错误', f'预测过程中出现错误: {str(e)}')

class ModelComparisonDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.parent = parent
        self.setWindowTitle("模型性能比较")
        self.setGeometry(200, 200, 700, 500)  # 稍微增大窗口尺寸
        
        # 创建布局
        layout = QVBoxLayout(self)
        
        # 模型选择区域
        selection_layout = QGridLayout()
        
        # 第一个模型选择
        selection_layout.addWidget(QLabel("模型 1:"), 0, 0)
        self.model1_combo = QComboBox()
        self.model1_combo.addItems(['CNN', 'LSTM', 'RNN_RELU', 'ResNet18', 'ResNet34', 
                                    'ResNet50', 'ResNet101', 'ResNet152', 'DenseNet'])
        selection_layout.addWidget(self.model1_combo, 0, 1)
        
        # 第二个模型选择
        selection_layout.addWidget(QLabel("模型 2:"), 0, 2)
        self.model2_combo = QComboBox()
        self.model2_combo.addItems(['CNN', 'LSTM', 'RNN_RELU', 'ResNet18', 'ResNet34', 
                                    'ResNet50', 'ResNet101', 'ResNet152', 'DenseNet'])
        self.model2_combo.setCurrentIndex(1)  # 默认选择第二个模型
        selection_layout.addWidget(self.model2_combo, 0, 3)
        
        # 股票选择
        selection_layout.addWidget(QLabel("股票:"), 1, 0)
        self.stock_combo = QComboBox()
        for code in parent.dataset.stocks_codes[:20]:
            formatted_code = str(code).strip()[-6:]
            name = parent.stock_names.get(formatted_code, '未知股票')
            self.stock_combo.addItem(f'{formatted_code} - {name}', formatted_code)
        selection_layout.addWidget(self.stock_combo, 1, 1, 1, 3)
        
        # 比较按钮
        compare_btn = QPushButton("比较")
        compare_btn.clicked.connect(self.compare_models)
        selection_layout.addWidget(compare_btn, 2, 0, 1, 4)
        
        # 添加选择区域到主布局
        layout.addLayout(selection_layout)
        
        # 创建表格显示结果
        self.result_table = QTableWidget(3, 3)
        self.result_table.setHorizontalHeaderLabels(["评估指标", "模型1", "模型2"])
        self.result_table.setVerticalHeaderLabels(["", "", ""])  # 清除垂直表头
        self.result_table.horizontalHeader().setStretchLastSection(True)
        
        # 填充表格的初始值和指标名称
        metrics = ["均方误差(MSE)", "平均绝对误差(MAE)", "决定系数(R²)"]
        for i, metric in enumerate(metrics):
            self.result_table.setItem(i, 0, QTableWidgetItem(metric))
            for j in range(1, 3):
                self.result_table.setItem(i, j, QTableWidgetItem("-"))
        
        layout.addWidget(self.result_table)
        
        # 添加模型推荐结果标签
        self.recommendation_label = QLabel("请点击比较按钮获取模型推荐")
        self.recommendation_label.setStyleSheet("font-weight: bold; margin-top: 10px;")
        self.recommendation_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.recommendation_label)
    
    def compare_models(self):
        try:
            # 获取选择的模型和股票代码
            model1_name = self.model1_combo.currentText()
            model2_name = self.model2_combo.currentText()
            stock_code = self.stock_combo.currentData()
            
            # 股票代码格式处理
            pure_code = str(stock_code).strip()[-6:]
            
            # 寻找正确的文件路径
            possible_paths = [
                f'./data/stocks/{pure_code}.csv',
                f'./data/stocks/sh.{pure_code}.csv',
                f'./data/stocks/sz.{pure_code}.csv'
            ]
            
            df = None
            actual_path = None
            
            for path in possible_paths:
                if os.path.exists(path):
                    try:
                        df = pd.read_csv(path)
                        actual_path = path
                        break
                    except Exception as e:
                        print(f"尝试读取 {path} 失败: {str(e)}")
            
            if df is None:
                raise FileNotFoundError(f"找不到股票代码为 {pure_code} 的数据文件")
            
            # 确定用于预测的股票代码格式
            prediction_stock_code = pure_code
            if 'sh.' in actual_path:
                prediction_stock_code = f'sh.{pure_code}'
            elif 'sz.' in actual_path:
                prediction_stock_code = f'sz.{pure_code}'
            
            # 清理列名
            df.columns = df.columns.str.strip().str.lower()
            
            # 确保必要的列存在
            if 'close' not in df.columns:
                raise ValueError("数据文件中缺少 'close' 列，请检查文件格式。")
            
            # 提取收盘价
            true_close = df['close'].values
            
            # 定义收盘价预测函数
            def close_p(x):
                if type(x) == int:
                    return x
                x = x[0, 1].item()
                return x if 0.2 > x > -0.2 else 0.0
            
            # 获取两个模型的预测结果
            predict_func1 = getattr(self.parent.prediction, f'predict_{model1_name.lower()}')
            predict_close1 = [true_close[j] * (1 + close_p(predict_func1(prediction_stock_code, (0, j))))
                           for j in range(len(true_close))]
            
            predict_func2 = getattr(self.parent.prediction, f'predict_{model2_name.lower()}')
            predict_close2 = [true_close[j] * (1 + close_p(predict_func2(prediction_stock_code, (0, j))))
                           for j in range(len(true_close))]
            
            # 计算评估指标
            mse1 = mean_squared_error(true_close, predict_close1)
            mae1 = mean_absolute_error(true_close, predict_close1)
            r2_1 = r2_score(true_close, predict_close1)
            
            mse2 = mean_squared_error(true_close, predict_close2)
            mae2 = mean_absolute_error(true_close, predict_close2)
            r2_2 = r2_score(true_close, predict_close2)
            
            # 更新表格
            self.result_table.setItem(0, 1, QTableWidgetItem(f"{mse1:.4f}"))
            self.result_table.setItem(0, 2, QTableWidgetItem(f"{mse2:.4f}"))
            
            self.result_table.setItem(1, 1, QTableWidgetItem(f"{mae1:.4f}"))
            self.result_table.setItem(1, 2, QTableWidgetItem(f"{mae2:.4f}"))
            
            self.result_table.setItem(2, 1, QTableWidgetItem(f"{r2_1:.4f}"))
            self.result_table.setItem(2, 2, QTableWidgetItem(f"{r2_2:.4f}"))
            
            # 设置表头
            self.result_table.setHorizontalHeaderLabels(["评估指标", model1_name, model2_name])
            
            # 高亮更好的结果
            model1_score = 0
            model2_score = 0
            
            for i in [(0, mse1, mse2, False), (1, mae1, mae2, False), (2, r2_1, r2_2, True)]:
                row, val1, val2, higher_better = i
                if higher_better:
                    better_cell = 1 if val1 > val2 else 2
                    if val1 > val2:
                        model1_score += 1
                    else:
                        model2_score += 1
                else:
                    better_cell = 1 if val1 < val2 else 2
                    if val1 < val2:
                        model1_score += 1
                    else:
                        model2_score += 1
                
                self.result_table.item(row, better_cell).setBackground(Qt.green)
            
            # 评估并推荐模型
            stock_name = self.parent.stock_names.get(pure_code, pure_code)
            
            # 计算性能差异百分比
            mse_diff = abs((mse1 - mse2) / max(mse1, mse2)) * 100
            mae_diff = abs((mae1 - mae2) / max(mae1, mae2)) * 100
            r2_diff = abs(r2_1 - r2_2) * 100  # R²直接取差值的百分比
            
            avg_diff = (mse_diff + mae_diff + r2_diff) / 3
            
            # 根据得分和指标差异提供推荐
            if model1_score > model2_score:
                if avg_diff > 20:
                    recommendation = f"<b>推荐使用 {model1_name} 模型</b>，其性能明显优于 {model2_name} (平均优势: {avg_diff:.2f}%)"
                else:
                    recommendation = f"<b>轻微推荐 {model1_name} 模型</b>，其性能略优于 {model2_name} (平均优势: {avg_diff:.2f}%)"
            elif model2_score > model1_score:
                if avg_diff > 20:
                    recommendation = f"<b>推荐使用 {model2_name} 模型</b>，其性能明显优于 {model1_name} (平均优势: {avg_diff:.2f}%)"
                else:
                    recommendation = f"<b>轻微推荐 {model2_name} 模型</b>，其性能略优于 {model1_name} (平均优势: {avg_diff:.2f}%)"
            else:
                recommendation = f"<b>两个模型性能相当</b>，可以任选 {model1_name} 或 {model2_name}"
            
            self.recommendation_label.setText(
                f"对于股票 {stock_name} ({pure_code}):<br>{recommendation}"
            )
        
        except Exception as e:
            import traceback
            traceback.print_exc()
            QMessageBox.critical(self, '错误', f'比较模型时出错: {str(e)}')

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = PredictionGUI()
    window.show()
    sys.exit(app.exec_())